# -*- coding: utf-8 -*-

# @Author  : Saline
# @Time    : 2018/7/27 17:00
# @desc    :

from PIL import Image
from aip import AipOcr
import io,json
import base64
from colorama import init,Fore
from methods import get_baidu

import sys
default_encoding = 'utf-8'
if sys.getdefaultencoding() != default_encoding:
    reload(sys)
    sys.setdefaultencoding(default_encoding)

def ocr_img_baidu(image):
    # 百度OCR API  ，在 https://cloud.baidu.com/product/ocr 上注册新建应用即可
    """ 你的 APPID AK SK """

    client = AipOcr(APP_ID, API_KEY, SECRET_KEY)

    global combine_region
    # 切割题目+选项区域，左上角坐标和右下角坐标,自行测试分辨率
    combine_region = "10, 250, 1050, 1400".replace(' ','').split(',')
    combine_region = list(map(int, combine_region))
    region_im = image.crop((combine_region[0], combine_region[1], combine_region[2], combine_region[3]))

    img_byte_arr = io.BytesIO()
    region_im.save(img_byte_arr, format='PNG')

    image_data = img_byte_arr.getvalue()

    #with open('./9.png','a') as f:
    #    f.write(image_data)
    
    response = client.basicGeneral(image_data)
    words_result = response['words_result']

    texts = [x['words'] for x in words_result]

    if len(texts) > 2:
        question = texts[0]
        choices = texts[1:]
        choices = [x.replace(' ', '') for x in choices]
    else:
        print(Fore.RED + '截图区域设置错误，请重新设置' + Fore.RESET)
        exit(0)

    if len(choices)>2:

        # 处理出现问题为两行或三行
        if str(choices[1]).find("A:") !=-1:
            question += choices[0]
            choices.pop(0)
        elif "A:" in str(choices[2]):
            question += choices[0]
            question += choices[1]
            choices.pop(0)
            choices.pop(0)

    answer = ""
    for x in choices:
        answer += x +"\n"

    print("获取的问题:")
    print(question)
    print("提供的选项:")
    print(answer)

    keyword = str(question.replace('·',''))
    keyword += str(answer.strip('\n')).replace("\n"," ")
    print("参考页面")
    get_baidu(keyword)
    return question.replace('·',''), answer.strip('\n')


if __name__ == '__main__':
    #check_screenshot()
    image = Image.open("./screenshot.png")
    # question, choices = ocr_img(image)
    # print("Tess 识别结果:")
    # print(question)
    # print(choices)
    # print()
    question, choices = ocr_img_baidu(image)